test(spanner): unflake testCreateSessionDeadlineExceededWithNoSessionCreateWaitTime - #14866

          @Override
          public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount) {
            SpannerException spannerException = SpannerExceptionFactory.asSpannerException(t);
            if (MultiplexedSessionDatabaseClient.this.resourceNotFoundException.get() == null
                && (spannerException instanceof DatabaseNotFoundException
                    || spannerException instanceof InstanceNotFoundException
                    || spannerException instanceof SessionNotFoundException)) {
              MultiplexedSessionDatabaseClient.this.resourceNotFoundException.set(
                  (ResourceNotFoundException) spannerException);
            }

            SettableApiFuture<SessionReference> retryFuture = null;
            if (remainingAttempts > 1
                && !isClientClosed()
                && RETRYABLE_ERROR_CODES.contains(spannerException.getErrorCode())) {
              final SettableApiFuture<SessionReference> nextFuture = SettableApiFuture.create();
              if (MultiplexedSessionDatabaseClient.this.multiplexedSessionReference.compareAndSet(
                  sessionReferenceFuture, nextFuture)) {
                retryFuture = nextFuture;
              }
            }

            // 1. Unblock all waiters of the failed attempt first.
            // This is executed OUTSIDE any lock to avoid deadlocks in foreign listener callbacks.
            // Because the retryFuture was installed into multiplexedSessionReference above,
            // any waking waiter that retries immediately is guaranteed to observe the retry future.
            sessionReferenceFuture.setException(t);

            // 2. If a retry was scheduled, initiate it:
            if (retryFuture != null) {
              // Guard against concurrent client close that happened while waking waiters.
              if (isClientClosed()) {
                retryFuture.setException(newClosedException());
                return;
              }
              try {
                asyncCreateMultiplexedSession(retryFuture, remainingAttempts - 1);
              } catch (Throwable e) {
                // Failsafe: ensure waiting threads never hang if RPC scheduling throws synchronously.
                retryFuture.setException(e);
              }
            }
          }

  void close() {
    boolean releaseChannelUsage = false;
    ApiFuture<SessionReference> currentSessionFuture = null;
    synchronized (this) {
      if (!this.isClosed) {
        this.isClosed = true;
        this.maintainer.stop();
        releaseChannelUsage = true;
        currentSessionFuture = this.multiplexedSessionReference.get();
      }
    }
    if (releaseChannelUsage) {
      // If a session creation or retry is currently in flight, fail it immediately so
      // transactions waiting on the closed client do not hang.
      if (currentSessionFuture instanceof SettableApiFuture && !currentSessionFuture.isDone()) {
        ((SettableApiFuture<SessionReference>) currentSessionFuture)
            .setException(newClosedException());
      }
      spanner.getRpc().unregisterChannelPrimeSessionSource(this);
      synchronized (CHANNEL_USAGE) {
        SharedChannelUsage sharedChannelUsage = CHANNEL_USAGE.get(this.spanner);
        if (sharedChannelUsage != null) {
          sharedChannelUsage.referenceCount--;
          if (sharedChannelUsage.referenceCount == 0) {
            CHANNEL_USAGE.remove(this.spanner);
          }
        }
      }
    }
  }

  @VisibleForTesting
  ApiFuture<SessionReference> getMultiplexedSessionFuture() {
    return this.multiplexedSessionReference.get();
  }

  @Test
  public void testRetryOfInitialSessionIsInstalledBeforeFailureIsPropagated() throws Exception {
    Clock clock = mock(Clock.class);
    when(clock.instant()).thenReturn(Instant.now());
    SessionClient sessionClient = mock(SessionClient.class);
    SpannerImpl spanner = mock(SpannerImpl.class);
    SpannerOptions spannerOptions = mock(SpannerOptions.class);
    SessionPoolOptions sessionPoolOptions = mock(SessionPoolOptions.class);
    when(sessionClient.getSpanner()).thenReturn(spanner);
    when(sessionClient.getDatabaseId()).thenReturn(TEST_DATABASE_ID);
    when(spanner.getRpc()).thenReturn(mock(SpannerRpc.class));
    when(spanner.getOptions()).thenReturn(spannerOptions);
    when(spannerOptions.getSessionPoolOptions()).thenReturn(sessionPoolOptions);
    when(sessionPoolOptions.getMultiplexedSessionMaintenanceDuration())
        .thenReturn(Duration.ofDays(7));
    when(sessionPoolOptions.getMultiplexedSessionMaintenanceLoopFrequency())
        .thenReturn(Duration.ofMinutes(10));
    SessionImpl session = mock(SessionImpl.class);
    SessionReference sessionReference = mock(SessionReference.class);
    when(session.getSessionReference()).thenReturn(sessionReference);

    AtomicReference<SessionConsumer> initialConsumer = new AtomicReference<>();
    AtomicReference<SessionConsumer> retryConsumer = new AtomicReference<>();
    doAnswer(
            (Answer<?>)
                invocationOnMock -> {
                  initialConsumer.set(invocationOnMock.getArgument(0));
                  return null;
                })
        .doAnswer(
            (Answer<?>)
                invocationOnMock -> {
                  retryConsumer.set(invocationOnMock.getArgument(0));
                  return null;
                })
        .when(sessionClient)
        .asyncCreateMultiplexedSession(any(SessionConsumer.class));

    MultiplexedSessionDatabaseClient client =
        new MultiplexedSessionDatabaseClient(sessionClient, clock);
    try {
      assertNotNull(initialConsumer.get());

      ApiFuture<SessionReference> initialFuture = client.getMultiplexedSessionFuture();
      AtomicReference<ApiFuture<SessionReference>> futureSeenByNewTransaction =
          new AtomicReference<>();
      initialFuture.addListener(
          () -> futureSeenByNewTransaction.set(client.getMultiplexedSessionFuture()),
          MoreExecutors.directExecutor());

      initialConsumer
          .get()
          .onSessionCreateFailure(
              SpannerExceptionFactory.newSpannerException(ErrorCode.DEADLINE_EXCEEDED, "test"), 1);

      // 1. Waiters for the initial attempt receive the error of that attempt.
      SpannerException exception =
          assertThrows(SpannerException.class, () -> SpannerApiFutures.get(initialFuture));
      assertEquals(ErrorCode.DEADLINE_EXCEEDED, exception.getErrorCode());

      // 2. New transactions arriving after the failure observe the retry future.
      assertNotNull(futureSeenByNewTransaction.get());
      assertFalse(futureSeenByNewTransaction.get().isDone());
      assertNotNull(retryConsumer.get());

      // 3. Completing the retry succeeds the pending future.
      retryConsumer.get().onSessionReady(session);
      assertEquals(sessionReference, futureSeenByNewTransaction.get().get());
      assertEquals(sessionReference, client.getCurrentSessionReference());
    } finally {
      client.close();
    }
  }